"""8.2 实现集成外部mcp服务npx类型"""
import asyncio
from contextlib import AsyncExitStack
from typing import Optional

from mcp import StdioServerParameters, ClientSession, ListToolsResult, ListResourcesResult
from mcp.client.stdio import stdio_client
from openai import OpenAI


class MCPClient:

    server_path= "server.py"

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()

    """创建连接服务端"""
    async def connect_to_server(self, server_path: str):
        # 一、创建StdioServerParameter参数信息
        server_parameters = StdioServerParameters(
            command="uvx",
            args=[  "mcp-server-sqlite",
                    "--db-path",
                    "D:\\dir_test\\db.sqlite3"
                 ],
            env=None
        )
        # 二、 创建stdioClient
        client = stdio_client(server=server_parameters)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession
        client_session = ClientSession(read_stream, write_stream)
        self.session:ClientSession = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()
        #五、列出所有可用的工具
        list_tool:ListToolsResult = await self.session.list_tools()
        tools = list_tool.tools
        for tool in tools:
           print("tool ==> ",tool.name,",",tool.description,",",tool.inputSchema)

        #列出所有的资源
        list_resources_result:ListResourcesResult = await self.session.list_resources()
        resources = list_resources_result.resources
        for resource in resources:
            print("resource ==> ",resource.name,",",resource.uri,",",resource.description,",",resource.mimeType)

        # 列表所有提示类
        list_prompts_result = await self.session.list_prompts()
        prompts = list_prompts_result.prompts
        for prompt in prompts:
            print("prompt ==> ",prompt.name,",",prompt.description,",",prompt.arguments)
    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server(client.server_path)
    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await client.cleanup()
        # 添加一小段延迟确保资源完全释放
        await asyncio.sleep(0.1)

if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(main())
    finally:
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.close()